from tree.multi_trees import *
from tree.mcts import MCTS
from opt_func.FuncSet import *
from opt_func.FuncApi import Function
import time
from threading import Thread
import sys
sys.path.append('./models/')
# 按间距中的绿色按钮以运行脚本。
if __name__ == '__main__':
    # 设置随机数种子
    np.random.seed(42)
    # 变量和取值范围的定义
    # x = np.round(np.linspace(-10,0,300),2)
    # y = np.round(np.linspace(-6.5,0,200),2)
    # values_range_list = [x,y]
    value_num = 10
    extra_num = 1
    low = -5
    high = 5
    split_num = int((high - low)/0.2) + 1

    tree_num, value_index = generate_multi_index(value_num, extra_num)
    values_range_list = np.round(np.repeat(np.linspace(
        low, high, split_num), value_num).reshape(split_num, -1).T, 2)
    function = Rosenbrock
    print("开始训练")
    flag = 1
    if flag:
        print("单数训练")
        # # 单课树的优化
        fn = Function(data_index=value_index[0], function=function,
                      values_range=values_range_list)
        mcts = MCTS(fn, n_playout=1*int(1e4), c_values=0.5e3, simu_times=10000,
                    pause_thre=0)

        mcts.run(select_rate=[0, 9, 0.7, 0.6])

        print("-------------------------")
        print("输出结果")
        print("最优序列组合为：")
        mcts.output_the_best()
        print("-------------------------")
        print("各种阈值下树的变量过滤情况")
        v_d = {}
        for thre in [0.3, 0.2, 0.1, 0.05]:
            v_d[thre] = mcts.merge_layer(thre)
        print(v_d)
    else:
        print("现在是多树训练")
        # 多棵树的训练
        mcts_list = []
        t_list = []
        for i in range(tree_num):
            fn = Function(data_index=value_index[i], function=function,
                          values_range=values_range_list)
            # mcts = MCTS(fn,n_playout=40000,c_values=0.5e3,simu_times=10000,
            #             pause_thre=0,pause_times=10000,stop_layer_play_time=8000,extra_layer=extra_num)
            mcts = MCTS(fn, n_playout=40000, c_values=0.5*1e3, simu_times=10000,
                        pause_thre=0)
            mcts_list.append(mcts)
            t = Thread(target=mcts.run, args=([1, 0.8],))
            t_list.append(t)
        t1 = time.time()
        for t in t_list:
            t.start()
        for t in t_list:
            t.join()

        print("输出结果")
        print(f"总运行时间为{time.time() - t1}s")
        print("输出每棵树的优化结果")
        for i, mctc_ in enumerate(mcts_list):
            print("-------------------------")
            print("第%d棵树的优化结果为：" % (i+1))
            mctc_.output_the_best()
            print("-------------------------")
        # TODO:对于多树的训练，可以用其他树来判定彼此训练的效果，如何对于表现不好的要求他重新训练（带指导的）
        print("过滤出的多树集合为")
        v_d = {}
        for thre in [0.3, 0.2, 0.1, 0.05]:
            v_d[thre] = extra_layers(
                mcts_list, extra_num, np.full(tree_num, thre))
        print(v_d)
